local math = require "math"
local table = require "table"

--[[
	Create a new Binary Search Tree with the given values. Be aware of the
	consequences of inserting already ordered data before you do this.
]]
local bst = {}
function bst:new(o)
	o = o or {}
	setmetatable(o, self)
	self.__index = self
	return o
end

function bst:new_multi(values)
	local new = self:new()
	for i, v in ipairs(values) do
		new:insert(v)
	end
	return new
end

--[[
	Insert a new value into a Binary Search tree. If replace evaluates to true,
	replace any existing equal value with the new one. Returns true or false
	to indicate if a replacement happened, and the node that was replaced or
	inserted
]]
function bst:insert(value, replace)
	if not self.value then
		self.value = value
		return false, self
	end
	if self.value == value then
		if replace then
			self.value = value
			return true, self
		end
	elseif self.value > value then
		if not self.left then
			self.left = bst:new{parent=self, value=value}
			return false, self.left
		else
			return self.left:insert(value, replace)
		end
	else
		if not self.right then
			self.right = bst:new{parent=self, value=value}
			return false, self.right
		else
			return self.right:insert(value, replace)
		end
	end
end

function bst:insert_multi(values, replace)
	for i, v in ipairs(values) do
		self:insert(v, replace)
	end
end

--[[
	Given the root node of a tree, find the smallest node on the tree. Return
	that node and it's value.
]]
function bst.first(node)
	while node.left do
		node = node.left
	end
	return node, node.value
end

--[[
	Given the root node of a tree, find the largest node on the tree. Return
	that node and it's value.
]]
function bst.last(node)
	while node.right do
		node = node.right
	end
	return node, node.value
end

--[[
	Return the next larger node and its value in the whole tree.
]]
function bst.next(node)
	if node.right then
		return node.right:first()
	end
	while node.parent do
		if node.parent.left == node then
			return node.parent, node.parent.value
		end
		node = node.parent
	end
end

--[[
	Return the next smaller node and its value in the whole tree
]]
function bst.previous(node)
	if node.left then
		return node.left:last()
	end
	while node.parent do
		if node.parent.right == node then
			return node.parent, node.parent.value
		end
		node = node.parent
	end
end

--[[
	Use in a for loop to iterate through all the nodes in a Binary Search
	Tree in sorted order.
]]
function bst:iter()
	local iter_function = self.first
	local node, value = self, nil
	return function ()
		node, value = iter_function(node)
		iter_function = self.next
		return node, value
	end
end

--[[
	Use in a for loop to iterate through all the nodes in a Binary Search
	Tree in reverse sorted order.
]]
function bst:rev_iter()
	local iter_function = self.last
	local node, value = self, nil
	return function ()
		node, value = iter_function(node)
		iter_function = self.previous
		return node, value
	end
end

--[[
	Determine whether a value already exists in a Binary Search Tree. Return
	the containing node and the value if it does. Else, returns false.
]]
function bst.search(node, value)
	if not node then
		return false
	elseif node.value > value then
		return bst.search(node.left, value)
	elseif node.value < value then
		return bst.search(node.right, value)
	end
	return node, node.value
end

--[[
	Return the number of direct children that node has. This will be either 0,
	1, or 2.
]]
local function count_children(node)
	local count = 0
	if node.left then
		count = count + 1
	end
	if node.right then
		count = count + 1
	end
	return count
end

--[[
	Delete a node with 0 or 1 direct child.
]]
local function bst_delete_simple(node)
	local replacement = node.left or node.right
	if node == node.parent.right then
		node.parent.right = replacement
	else
		node.parent.left = replacement
	end
	if replacement then
		replacement.parent = node.parent
	end
	node.parent = nil
end

--[[
	Delete a node with 2 direct children.
]]
local function bst_delete_complex(node)
	local replacement = nil
	if math.random(2) == 1 then
		replacement = node:next()
	else
		replacement = node:previous(node)
	end
	node.value = replacement.value
	bst_delete_simple(replacement)
end

--[[
	Delete a node with the given value from the tree. Returns true if the
	deletion occurred.
]]
function bst:delete(value)
	local to_delete = self:search(value)
	if not to_delete then
		return false
	end
	if count_children(to_delete) == 2 then
		bst_delete_complex(to_delete)
	else
		bst_delete_simple(to_delete)
	end
	return true
end

function bst:delete_multi(values)
	for i, v in ipairs(values) do
		self:delete(v)
	end
end

function bst:__tostring()
	local queue = {}
	local results = {}
	table.insert(queue, self)
	while queue[1] do
		local dequeued = table.remove(queue, 1)
		table.insert(results, dequeued.value)
		if dequeued.left then
			table.insert(queue, dequeued.left)
		end
		if dequeued.right then
			table.insert(queue, dequeued.right)
		end
	end
	local as_strings = {}
	for i, v in ipairs(results) do
		table.insert(as_strings, tostring(v))
	end
	return table.concat(as_strings, " ")
end

return bst

